import torch
from ultralytics import YOLO

# 设备配置
device = 0 if torch.cuda.is_available() else "cpu"

# ================= 第一阶段训练 =================
# 初始化模型（加载预训练权重）
model = YOLO("yolo11m.pt")

# 第一阶段训练配置
phase1_params = {
    "data": "dataset.yaml",
    "epochs": 30,                # 总epochs的30%
    "freeze": [0, 1, 2],        # 冻结前3个stage
    "imgsz": 1024,
    "batch": 8,
    "device": device,
    "lr0": 0.001,               # 初始学习率
    "cos_lr": True,             # 启用余弦退火
    "warmup_epochs": 5,         # 学习率预热
    "optimizer": "AdamW",
    "weight_decay": 0.05,
    "patience": 15,             # 早停观察周期
    "mosaic": 0.7,              # 保持较高mosaic概率
    "degrees": 5.0,             # 较小幅度旋转
    "name": "phase1_frozen"     # 保存路径标识
}

# 执行第一阶段训练
phase1_results = model.train(**phase1_params)
# ================= 第二阶段训练 =================
# 加载第一阶段最佳权重
model = YOLO("runs/detect/phase1_frozen/weights/best.pt")

# 第二阶段训练配置
phase2_params = {
    "data": "dataset.yaml",
    "epochs": 70,               # 总epochs的70%
    "imgsz": 1024,
    "batch": 8,
    "device": device,
    "lr0": 0.0001,              # 更小的学习率
    "cos_lr": True,
    "warmup_epochs": 3,         # 缩短预热时间
    "optimizer": "AdamW",
    "weight_decay": 0.01,        # 降低权重衰减
    "patience": 20,
    "mosaic": 0.3,              # 降低mosaic概率
    "copy_paste": 0.2,           # 启用复制粘贴增强
    "degrees": 10.0,             # 增强幅度增大
    "dropout": 0.2,              # 添加Dropout
    "close_mosaic": 10,          # 最后10轮关闭增强
    "name": "phase2_finetune"    # 新保存路径
}

# 执行第二阶段训练
phase2_results = model.train(**phase2_params)
# ================= 最终评估 =================
# 加载最终模型
final_model = YOLO("runs/detect/phase2_finetune/weights/best.pt")

# 统一验证配置
val_config = {
    "data": "dataset.yaml",
    "conf": 0.25,
    "iou": 0.45,
    "plots": True
}

# 验证集评估
val_metrics = final_model.val(**val_config)

# 测试集评估
test_metrics = final_model.val(
    **val_config,
    split="test"  # 使用测试集分割
)

# 模型导出（ONNX格式）
# export_path = final_model.export(
#     format="onnx",
#     simplify=True,
#     opset=12,
#     dynamic=False,
#     device=device
# )